import numpy as np
from scipy.signal import butter, sosfiltfilt

def extractEpoch3D(data, event, srate, baseline, frame, opt_keep_baseline):
    # extract epoch from 2D data into 3D [ch x time x trial]
    # time才是数据
    # input: event, baseline, frame
    # extract epoch = baseline[0] to frame[2]

    # for memory pre-allocation
    if opt_keep_baseline == True:
        begin_tmp = int(np.floor(baseline[0] / 1000 * srate))
        end_tmp = int(begin_tmp + np.floor(frame[1] - baseline[0]) / 1000 * srate)
    else:
        begin_tmp = int(np.floor(frame[0] / 1000 * srate))
        end_tmp = int(begin_tmp + np.floor(frame[1] - frame[0]) / 1000 * srate)
    # shape为信号、信号、第几次
    epoch3D = np.zeros((data.shape[0], end_tmp - begin_tmp, len(event)))
    nth_event = 0

    for i in event:
        if opt_keep_baseline == True:
            begin_id = int(i + np.floor(baseline[0] / 1000 * srate))
            end_id = int(begin_id + np.floor((frame[1] - baseline[0]) / 1000 * srate))
        else:
            begin_id = int(i + np.floor(frame[0] / 1000 * srate))
            end_id = int(begin_id + np.floor((frame[1] - frame[0]) / 1000 * srate))

        tmp_data = data[:, begin_id:end_id]

        begin_base = int(np.floor(baseline[0] / 1000 * srate))
        end_base = int(begin_base + np.floor(np.diff(baseline) / 1000 * srate) - 1)
        base = np.mean(tmp_data[:, begin_base:end_base], axis=1)

        rmbase_data = tmp_data - base[:, np.newaxis]
        epoch3D[:, :, nth_event] = rmbase_data
        nth_event = nth_event + 1

    return epoch3D


def butter_bandpass_filter(data, lowcut, highcut, fs, order):
    nyq = fs / 2
    low = lowcut / nyq
    high = highcut / nyq
    sos = butter(order, [low, high], btype='band', output='sos')
    # demean before filtering
    meandat = np.mean(data, axis=1)
    data = data - meandat[:, np.newaxis]
    y = sosfiltfilt(sos, data)  # zero-phase filter # data: [ch x time]
    # specify pandlen to make the result the same as Matlab filtfilt()
    return y